import torch
from PIL import Image
# from paddle import fluid
# import paddle
from sklearn import metrics
def create_loss(predict, label, num_classes=2):
    ''' 创建loss，结合dice和交叉熵 '''
    predict = fluid.layers.transpose(predict, perm=[0, 2, 3, 1])
    predict = fluid.layers.reshape(predict, shape=[-1, num_classes])
    predict = fluid.layers.softmax(predict)
    label = fluid.layers.reshape(label, shape=[-1, 1])
    label = fluid.layers.cast(label, "int64")
    dice_loss = fluid.layers.dice_loss(predict, label)  # 计算dice loss
    ce_loss = fluid.layers.cross_entropy(predict, label)  # 计算交叉熵
    return fluid.layers.reduce_mean(ce_loss + dice_loss)  # 最后使用的loss是dice和交叉熵的和，单独使用dice一般不是很稳定


def mean_iou(pred, label, num_classes=2):
    ''' 计算miou，评价网络分割结果的指标'''
    paddle.disable_static()
    pred =paddle.to_tensor(np.array(pred))
    label =paddle.to_tensor(np.array(label))
    # pred = fluid.layers.argmax(pred)

    pred = fluid.layers.cast(pred, 'int32')
    label = fluid.layers.cast(label, 'int32')
    miou, wrong, correct = fluid.layers.mean_iou(pred, label, num_classes)
    return miou


def imgScale(img, newSize):
    # 等比例缩放至size大小的图片
    img.thumbnail((newSize, newSize))
    # 创建一个size大小的背景
    newImg = Image.new('RGB', (newSize, newSize), (0, 0, 0))

    # 获取图片的大小
    w, h = img.size

    # 将前景图片贴到背景图片上
    # newImg.paste(img, (0, 0))
    if w == newSize:
        newImg.paste(img, (0, int((newSize - h) / 2)))
    elif h == newSize:
        newImg.paste(img, (int((newSize - w) / 2), 0))
    else:
        newImg.paste(img, (int((newSize - w) / 2), int((newSize - h) / 2)))

    return newImg


import numpy as np

__all__ = ['SegmentationMetric']

"""
confusionMetric  # 注意：此处横着代表预测值，竖着代表真实值，与之前介绍的相反
P\L     P    N
P      TP    FP
N      FN    TN
"""


class SegmentationMetric(object):
    def __init__(self, numClass):
        self.numClass = numClass
        self.confusionMatrix = np.zeros((self.numClass,) * 2)

    def pixelAccuracy(self):
        # return all class overall pixel accuracy
        #  PA = acc = (TP + TN) / (TP + TN + FP + TN)
        acc = np.diag(self.confusionMatrix).sum() / self.confusionMatrix.sum()
        return acc

    def classPixelAccuracy(self):
        # return each category pixel accuracy(A more accurate way to call it precision)
        # acc = (TP) / TP + FP
        classAcc = np.diag(self.confusionMatrix) / self.confusionMatrix.sum(axis=1)
        return classAcc  # 返回的是一个列表值，如：[0.90, 0.80, 0.96]，表示类别1 2 3各类别的预测准确率

    def meanPixelAccuracy(self):
        classAcc = self.classPixelAccuracy()
        meanAcc = np.nanmean(classAcc)  # np.nanmean 求平均值，nan表示遇到Nan类型，其值取为0
        return meanAcc  # 返回单个值，如：np.nanmean([0.90, 0.80, 0.96, nan, nan]) = (0.90 + 0.80 + 0.96） / 3 =  0.89

    def meanIntersectionOverUnion(self):
        # Intersection = TP Union = TP + FP + FN
        # IoU = TP / (TP + FP + FN)
        intersection = np.diag(self.confusionMatrix)  # 取对角元素的值，返回列表
        union = np.sum(self.confusionMatrix, axis=1) + np.sum(self.confusionMatrix, axis=0) - np.diag(
            self.confusionMatrix)  # axis = 1表示混淆矩阵行的值，返回列表； axis = 0表示取混淆矩阵列的值，返回列表
        IoU = intersection / union  # 返回列表，其值为各个类别的IoU
        mIoU = np.nanmean(IoU)  # 求各类别IoU的平均
        return mIoU

    def genConfusionMatrix(self, imgPredict, imgLabel):  # 同FCN中score.py的fast_hist()函数
        # remove classes from unlabeled pixels in gt image and predict
        mask = (imgLabel >= 0) & (imgLabel < self.numClass)
        label = self.numClass * imgLabel[mask] + imgPredict[mask]
        count = np.bincount(label, minlength=self.numClass ** 2)
        confusionMatrix = count.reshape(self.numClass, self.numClass)
        return confusionMatrix

    def Frequency_Weighted_Intersection_over_Union(self):
        # FWIOU =     [(TP+FN)/(TP+FP+TN+FN)] *[TP / (TP + FP + FN)]
        freq = np.sum(self.confusion_matrix, axis=1) / np.sum(self.confusion_matrix)
        iu = np.diag(self.confusion_matrix) / (
                np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -
                np.diag(self.confusion_matrix))
        FWIoU = (freq[freq > 0] * iu[freq > 0]).sum()
        return FWIoU

    def addBatch(self, imgPredict, imgLabel):
        assert imgPredict.shape == imgLabel.shape
        self.confusionMatrix += self.genConfusionMatrix(imgPredict, imgLabel)

    def reset(self):
        self.confusionMatrix = np.zeros((self.numClass, self.numClass))


if __name__ == '__main__':
    imgPredict = np.array([0, 0, 1, 1, 2, 2])  # 可直接换成预测图片
    imgLabel = np.array([0, 0, 1, 1, 2, 2])  # 可直接换成标注图片
    metric = SegmentationMetric(1)  # 3表示有3个分类，有几个分类就填几
    metric.addBatch(imgPredict, imgLabel)
    pa = metric.pixelAccuracy()
    cpa = metric.classPixelAccuracy()
    mpa = metric.meanPixelAccuracy()
    mIoU = metric.meanIntersectionOverUnion()

# if __name__ == '__main__':
#     testFile = './VOC2012/JPEGImages/2007_000032.jpg'
#
#     img = Image.open(testFile)
#
#     img = imgScale(img, 240)
#
#     img.show()
